import os
from PIL import Image
from torch.utils.data import Dataset
import torch.nn.functional as F

class LeapGestRecogData(Dataset):
    def __init__(self, root, transform=None):
        self.root = root
        self.transform = transform
        self.x = []
        self.y = []

        folders = os.listdir(root)
        for folder in folders:
            for dirpath, _, filenames in os.walk(os.path.join(root, folder)):
                for filename in filenames:
                    class_name = os.path.basename(dirpath)
                    class_id = int(class_name.split("_")[0]) - 1
                    if class_id < 6:
                        self.x.append(os.path.join(dirpath, filename))
                        self.y.append(class_id)
        self.len = len(self.x)

    def __len__(self):
        return self.len

    def __getitem__(self, index):
        img = Image.open(self.x[index]).convert('L')
        y = self.y[index]
        if self.transform:
            img = self.transform(img)
        return img, y